from pymongo import MongoClient
from typing import List, Dict, Any, Optional
from config import config
import json

class DatabaseManager:
    """数据库管理模块"""
    
    def __init__(self):
        self.client = None
        self.db = None
        self.collection = None
        self.connect()
    
    def connect(self):
        """连接MongoDB数据库"""
        try:
            print("🔌 连接MongoDB...")
            self.client = MongoClient(config.MONGODB_URI, serverSelectionTimeoutMS=3000)
            self.client.admin.command('ping')
            print("MongoDB 连接成功")
            
            self.db = self.client[config.DB_NAME]
            self.collection = self.db[config.COLLECTION_NAME]
            print(f"使用数据库: {config.DB_NAME}, 集合: {config.COLLECTION_NAME}")
            
        except Exception as e:
            print(f"MongoDB连接失败: {str(e)}")
            raise
    
    def clear_collection(self):
        """清空集合"""
        try:
            self.collection.delete_many({})
            print("🗑 已清空集合")
        except Exception as e:
            print(f"清空集合失败: {str(e)}")
            raise
    
    def insert_documents(self, documents: List[Dict[str, Any]]):
        """插入文档到数据库"""
        try:
            # 转换numpy数组为列表以便JSON序列化
            for doc in documents:
                if "embedding" in doc:
                    doc["embedding"] = doc["embedding"].tolist() if hasattr(doc["embedding"], "tolist") else doc["embedding"]
            
            result = self.collection.insert_many(documents)
            print(f"✅ 成功插入 {len(result.inserted_ids)} 个文档")
            return result.inserted_ids
        except Exception as e:
            print(f"插入文档失败: {str(e)}")
            raise
    
    def get_all_documents(self) -> List[Dict[str, Any]]:
        """获取所有文档"""
        try:
            documents = list(self.collection.find({}, {"_id": 0}))
            print(f"📚 从数据库获取 {len(documents)} 个文档")
            return documents
        except Exception as e:
            print(f"获取文档失败: {str(e)}")
            raise
    
    def search_documents(self, query: Dict[str, Any]) -> List[Dict[str, Any]]:
        """搜索文档"""
        try:
            documents = list(self.collection.find(query, {"_id": 0}))
            return documents
        except Exception as e:
            print(f"搜索文档失败: {str(e)}")
            raise
    
    def update_document(self, filter_query: Dict[str, Any], update_data: Dict[str, Any]):
        """更新文档"""
        try:
            result = self.collection.update_one(filter_query, {"$set": update_data})
            return result.modified_count
        except Exception as e:
            print(f"更新文档失败: {str(e)}")
            raise
    
    def delete_document(self, filter_query: Dict[str, Any]):
        """删除文档"""
        try:
            result = self.collection.delete_one(filter_query)
            return result.deleted_count
        except Exception as e:
            print(f"删除文档失败: {str(e)}")
            raise
    
    def get_document_count(self) -> int:
        """获取文档数量"""
        try:
            return self.collection.count_documents({})
        except Exception as e:
            print(f"获取文档数量失败: {str(e)}")
            return 0
    
    def close(self):
        """关闭数据库连接"""
        if self.client:
            self.client.close()
            print("数据库连接已关闭")

# 全局数据库管理器实例
db_manager = DatabaseManager() 